#!/usr/bin/env python
# encoding: utf-8
"""
@author: jfjiang6
@file: ner_crf_sklearn_cfrsuite.py

"""
import pandas as pd
import pickle
from codecs import open
from collections import Counter

from sklearn_crfsuite import CRF
from sklearn_crfsuite.metrics import flat_precision_score, flat_recall_score, flat_f1_score

import warnings
warnings.filterwarnings('ignore')

# #############################################################################################
# todo: 构建类别标注的评价体系并存储在evaluation中，全局变量
evaluation = pd.DataFrame({'tag': [],
                           '精确率': [],
                           '召回率': [],
                           ' F1 值': [],
                           'support数目': [],
                           'predict数目': []})


# #############################################################################################
# 步骤1 todo: 定义一些辅助函数，用于特征提取
#  步骤1.1 todo: 定义函数抽取单个字的特征
def word2features(sent, i):  # sent 为一个句子单位
    """
    针对句子中的第i个单词，记录这个句子中，该单词的前后单词关系。
    :param sent:
    :param i:
    :return:
    """
    word = sent[i] # 取出句子中的某个单词
    prev_word = "<s>" if i == 0 else sent[i - 1]
    next_word = "</s>" if i == (len(sent) - 1) else sent[i + 1]
    # 使用的特征：
    # 前一个词，当前词，后一个词，
    # 前一个词+当前词， 当前词+后一个词
    features = {
        'w': word,
        'w-1': prev_word,
        'w+1': next_word,
        'w-1:w': prev_word + word,
        'w:w+1': word + next_word,
        'bias': 1
    }
    return features


# 步骤1.2 todo: 定义函数抽取序列特征
def sent2features(sent):
    """
    遍历这个句子中的所有单词，均记录单词的前后单词关系
    :param sent:
    :return: 每个单词对应的前后单词关系的列表。
    """
    return [word2features(sent, i) for i in range(len(sent))]


# 步骤1.3 todo: 定义函数将列表的列表转化为列表: [[t1, t2], [t3, t4]...] --> [t1, t2, t3, t4...]
def flatten_lists(lists):
    """
    将二维列表转换成一维列表
    :param lists:
    :return:
    """
    flatten_list = []
    for l in lists:
        if type(l) == list: # 判断l是否为列表list类型
            flatten_list += l
        else:
            flatten_list.append(l)
    return flatten_list


# #############################################################################################
# 步骤2 todo: 定义函数加载中文简历数据集
def load_corpus(filename):
    word_lists = [] # 简历相关的词
    tag_lists = [] # 词对应的标签
    with open(filename, 'r', encoding='utf-8') as f:
        word_list = []
        tag_list = []
        for line in f:
            if line != '\n':
                word, tag = line.strip('\n').split()
                word_list.append(word)
                tag_list.append(tag)
            else: # 换行符代表读完一个完整的句子，将这个句子的单词和每个单词对应的标签加入到word_lists和tag_lists中
                word_lists.append(word_list)
                tag_lists.append(tag_list)
                word_list = []
                tag_list = []
    return word_lists, tag_lists


# #############################################################################################
# 步骤3 todo: 定义函数用于CRF模型命名实体识别训练、预测和评价
def crf_ner(model, train_word_lists, train_tag_lists, test_word_lists, test_tag_lists):

    # 步骤3.1 todo: 分别对训练词列表和测试词列表进行特征抽取
    train_features = [sent2features(s) for s in train_word_lists] # 将每个句子中的每个词记录它的相邻词关系
    test_features = [sent2features(s) for s in test_word_lists]

    # 步骤3.2 todo: 模型训练
    crf_model = model
    crf_model.fit(train_features, train_tag_lists)

    # 步骤3.3 todo: 模型保存
    with open('crf_model.pkl', "wb") as model:
        pickle.dump(crf_model, model) # 将模型写入文件

    # 步骤3.4 todo: 模型预测
    pred_tag_lists = crf_model.predict(test_features)

    # 步骤3.5 todo: 模型评价(计算精确率、召回率和F1值）
    # 步骤3.5.1 todo: 计算总体的精确率、召回率和F1值（所有标注类别的加权平均）
    precision_total = flat_precision_score(test_tag_lists, pred_tag_lists, average='micro')
    recall_total = flat_recall_score(test_tag_lists, pred_tag_lists, average='micro')
    f1_scores_total = flat_f1_score(test_tag_lists, pred_tag_lists, average='micro')

    # 步骤3.5.2 todo: 计算每个标注类别的精确率、召回率和F1值
    # 步骤3.5.2.1 todo: 将标注数据列表进行格式转化(列表的列表转化为列表）
    test_tag_lists = flatten_lists(test_tag_lists)
    pred_tag_lists = flatten_lists(pred_tag_lists)

    # 步骤3.5.2.2 todo: 统计标注数据列表（测试集和预测所得集）中的标注类别（类别及对应的个数）
    predict_tags_counter = Counter(pred_tag_lists)
    test_tags_counter = Counter(test_tag_lists)

    # 步骤3.5.2.3 todo: 计算每个标注类别预测正确的个数，用于后面精确率以及召回率的计算
    correct_tags_number = {}
    for test_tag, predict_tag in zip(test_tag_lists, pred_tag_lists):
        if test_tag == predict_tag: # 预测正确
            if test_tag not in correct_tags_number:
                correct_tags_number[test_tag] = 1
            else:
                correct_tags_number[test_tag] += 1

    # 步骤3.5.2.4 todo: 计算每个标注类别的精确率、召回率和F1值
    precision_scores = {}  # 精确率
    recall_scores = {}  # 召回率
    f1_scores = {}  # F1值
    test_tag_set = set(test_tag_lists)  # 将列表转化为集合（去掉了重复的标注类别），计算效率高
    for tag in test_tag_set:
        precision_scores[tag] = correct_tags_number.get(tag, 0) / predict_tags_counter[tag] # 预测为tag标签正确的样本数/预测为tag标签总样本数
        recall_scores[tag] = correct_tags_number.get(tag, 0) / test_tags_counter[tag] # 预测为tag标签正确的样本数/实际标签为tag的总样本数
        f1_scores[tag] = 2 * precision_scores[tag] * recall_scores[tag] / \
                         (precision_scores[tag] + recall_scores[tag] + 1e-10)  # 分母加上很小的数，防止为0

    # 步骤3.6 todo: 将评价结果存入evaluation中
    # 步骤3.6.1 todo: 将每个标注类别的精确率、召回率和F1值及对应的标注个数存入evaluation中
    for tag in test_tag_set:
        r = evaluation.shape[0]
        evaluation.loc[r] = [tag,
                             format(precision_scores[tag], '.5f'),
                             format(recall_scores[tag], '.5f'),
                             format(f1_scores[tag], '.5f'),
                             test_tags_counter[tag],
                             predict_tags_counter[tag]]

    # 步骤3.6.2 todo: 将总体的精确率、召回率和F1值存入evaluation中
    r = evaluation.shape[0]
    evaluation.loc[r] = ['avg/total',
                         format(precision_total, '.5f'),
                         format(recall_total, '.5f'),
                         format(f1_scores_total, '.5f'),
                         len(test_tag_lists),
                         len(pred_tag_lists)]

    # 步骤3.7 todo: 将评价指标evaluation保存，encoding防止中文乱码
    evaluation.to_csv('./evaluation.csv', sep=',', header=True, index=True, encoding='utf_8_sig')


# #############################################################################################
if __name__ == "__main__":
    # 步骤1 todo: 加载语料库，分别给出训练集和测试集
    print("加载语料库的训练集和测试集...")
    train_word_lists, train_tag_lists = load_corpus('./data/ResumeNER/train.char.bmes')
    test_word_lists, test_tag_lists = load_corpus('./data/ResumeNER/test.char.bmes')

    # 步骤2 todo: 给出模型实例
    # 设置CRF模型参数
    algorithm = 'lbfgs'  # 梯度下降选取L-bfgs算法
    c1 = 0.1  # L1惩罚项的系数
    c2 = 0.1  # L2惩罚项的系数
    max_iterations = 100  # 迭代次数
    all_possible_transitions = False # 描述相邻标签的依赖关系（如形容词后接名词的概率），False仅考虑训练数据中实际出现过的标签转移
    crf_model = CRF(algorithm=algorithm, c1=c1, c2=c2, max_iterations=max_iterations,
                    all_possible_transitions=all_possible_transitions)

    # 步骤3 todo: 进行CRF模型命名实体识别训练、测试及评价
    print("进行CRF模型训练、测试及评价...")
    crf_ner(crf_model, train_word_lists, train_tag_lists, test_word_lists, test_tag_lists)
    print("CRF模型训练、测试及评价结束！")
